#!/bin/sh
################################################################################
#
#   Copyright (C) 2009 Ivan Mironov <mironov.ivan@gmail.com>
#
#   This program is free software: you can redistribute it and/or modify
#   it under the terms of the GNU General Public License as published by
#   the Free Software Foundation, either version 3 of the License, or
#   (at your option) any later version.
#
#   This program is distributed in the hope that it will be useful,
#   but WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#   GNU General Public License for more details.
#
#   You should have received a copy of the GNU General Public License
#   along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
################################################################################

OPT_DEST="DEFAULT_DEST"
OPT_MODE=4750 # -rwsr-x---
OPT_MODE_A=4755 # -rwsr-xr-x

# Print error and exit.
# $1 - error message.
die()
{
	echo "ERROR: ${1}"
	[ "x${tmp}" != "x" ] && rm -rf "${tmp}"
	exit 1
}

# Convert string to escaped octal codes.
# $1 - string.
# stdout - escaped string.
strToOct()
{
	printf "%s" "${1}" | od -v -A n -t o1 | \
			sed "s/ *\([0-9]*\) */\\\\\1/g" | tr -d "\n"
}

# Generate wrapper source code.
# $1 - real executable.
genSource()
{
	local i argv envp

cat << EOF
/* For putenv(). */
#define _SVID_SOURCE
#define _XOPEN_SOURCE

#include <stdlib.h>
#include <stdio.h>
#include <errno.h>
#include <string.h>
#include <unistd.h>

EOF
	# File name of executable.
	printf "static const char *const real_exec = \""
	strToOct "${1}"
	printf "\";\n"
cat << EOF

static void printErr()
{
	fprintf(stderr, "\"%s\" wrapper error: %s (code %d)\n",
			real_exec, strerror(errno), errno);
	exit(EXIT_FAILURE);
}

int main(int argc, char *argv[])
{
EOF
	# For additions arguments and C90.
	if [ $OPT_ARG_NUM -gt 0 ]; then
		printf "\tint i;\n"
		printf "\tchar **new_argv;\n"
	fi

	if [ "x${OPT_EMPTY_ENV}" != "x" ]; then
		# Empty environment.
		printf "\tchar *new_envp[] = {NULL};\n"
		envp="new_envp"
	else
		printf "\textern char **environ;\n"

		# Set environment variables.
		for i in $( seq $OPT_ENV_NUM ); do
			printf "\tif (putenv(\""
			eval strToOct \"\${OPT_ENV_$i}\"
			printf "\"))\n"
			printf "\t\tprintErr();\n"
		done

		# Unset environment variables.
		for i in $( seq $OPT_UENV_NUM ); do
			printf "\tif (unsetenv(\""
			eval strToOct \"\${OPT_UENV_$i}\"
			printf "\"))\n"
			printf "\t\tprintErr();\n"
		done

		envp="environ"
	fi

	# Additions arguments.
	if [ $OPT_ARG_NUM -gt 0 ]; then
		printf "\tnew_argv = "
		printf "malloc(sizeof(*new_argv) * (argc + %s + 1));\n" \
				$OPT_ARG_NUM

		printf "\tnew_argv[0] = argv[0];\n"

		for i in $( seq $OPT_ARG_NUM ); do
			printf "\tnew_argv[%s] = \"" $i
			eval strToOct \"\${OPT_ARG_$i}\"
			printf "\";\n"
		done
		printf "\tfor (i = 1; i <= argc; ++i)\n"
		printf "\t\tnew_argv[i + %s] = argv[i];\n" $OPT_ARG_NUM

		argv="new_argv"
	else
		printf "\t(void) argc;\n"
		argv="argv"
	fi
cat << EOF
	execve(real_exec, ${argv}, ${envp});
	printErr();
	return 0; /* =) */
}
EOF
}

# Print usage and exit.
printUsage()
{
cat << EOF
Copyright (c) 2009 Ivan Mironov <mironov.ivan@gmail.com>
Usage:
	${0} [-hAEs] [-d destination-directory] [-m mode] [-o owner] [-g group]
		[-a argument] [-e ENVVAR=value] [-u ENVVAR] real-executable
		[wrapper-name]
EOF
	exit 2
}

# Verify that the utility is available. 
# $1 - name.
checkUtil()
{
	which "${1}" >/dev/null 2>&1 || die "\"${1}\" not available"
}

# Check required utilities.
checkUtil which
checkUtil rm
checkUtil od
checkUtil sed
checkUtil tr
checkUtil basename
checkUtil chmod
checkUtil install
checkUtil id
checkUtil cc

# Default owner and group.
OPT_OWNER=$( id -un )
OPT_GROUP=$( id -gn )

# Read command line options.
OPT_ARG_NUM=0
OPT_ENV_NUM=0
OPT_UENV_NUM=0
while getopts "hAEsd:m:o:g:a:e:u:" opt; do
	case "${opt}" in
		"h") printUsage;;
		"A") OPT_MODE="${OPT_MODE_A}";;
		"E") OPT_EMPTY_ENV="X";;
		"s") OPT_PRINT_SOURCE="X";;
		"d") OPT_DEST="${OPTARG}";;
		"m") OPT_MODE="${OPTARG}";;
		"o") OPT_OWNER="${OPTARG}";;
		"g") OPT_GROUP="${OPTARG}";;

		"a")
			OPT_ARG_NUM=$(( $OPT_ARG_NUM + 1 ))
			eval OPT_ARG_$OPT_ARG_NUM=\"\${OPTARG}\"
		;;

		"e")
			OPT_ENV_NUM=$(( $OPT_ENV_NUM + 1 ))
			eval OPT_ENV_$OPT_ENV_NUM=\"\${OPTARG}\"
		;;

		"u")
			OPT_UENV_NUM=$(( $OPT_UENV_NUM + 1 ))
			eval OPT_UENV_$OPT_UENV_NUM=\"\${OPTARG}\"
		;;

		*) die "Try ${0} -h";;
	esac
done
shift $(( $OPTIND - 1 ))
if [ $# -lt 1 ] || [ $# -gt 2 ]; then
	die "Invalid arguments, try ${0} -h"
fi
OPT_REAL="${1}"
OPT_WRAPPER="${2}"

[ "x${OPT_DEST}" = "x" ] && die "Empty destination path"
[ "x${OPT_MODE}" = "x" ] && die "Empty file mode"
[ "x${OPT_OWNER}" = "x" ] && die "Empty file owner"
[ "x${OPT_GROUP}" = "x" ] && die "Empty file group"
[ "x${OPT_REAL}" = "x" ] && die "Empty executable path"

# Get full path to script.
real_base=$( basename "${OPT_REAL}" )
[ "x${OPT_REAL}" = "x${real_base}" ] && OPT_REAL=$( which "${OPT_REAL}" )
[ "x${OPT_REAL}" = "x" ] && die "Executable \"${1}\" not found in \$PATH"

# Get full path to wrapper binary.
[ "x${OPT_WRAPPER}" = "x" ] && OPT_WRAPPER=$( basename "${OPT_REAL}" )
wrapper_base=$( basename "${OPT_WRAPPER}" )
[ "x${OPT_WRAPPER}" = "x${wrapper_base}" ] && \
		OPT_WRAPPER="${OPT_DEST}/${OPT_WRAPPER}"

if [ "x${OPT_PRINT_SOURCE}" != "x" ]; then
	# Print source and exit.
	genSource "${OPT_REAL}"
	exit 0
fi

[ -e "${OPT_WRAPPER}" ] && die "File \"${OPT_WRAPPER}\" already exists"

# Create directory for temporary files.
tmp=$( mktemp -d exec-wrapper.XXXXXXXXXX ) || \
		die "Can't create temporary directory"
chmod 0700 "${tmp}"

# Compile wrapper.
echo "Generating wrapper \"${OPT_WRAPPER}\" -> \"${OPT_REAL}\""
echo "Owner: \"${OPT_OWNER}\", group: \"${OPT_GROUP}\", mode: \"${OPT_MODE}\""

WRAPPER_SRC="${tmp}/wrapper.c"
WRAPPER_BIN="${tmp}/wrapper"
genSource "${OPT_REAL}" >"${WRAPPER_SRC}"
cc -o "${WRAPPER_BIN}" "${WRAPPER_SRC}" || die "Can't compile wrapper"

# Install wrapper.
install -o "${OPT_OWNER}" -g "${OPT_GROUP}" -m "${OPT_MODE}" -s \
		"${WRAPPER_BIN}" "${OPT_WRAPPER}" || \
		die "Can't install wrapper binary"

# Delete temporary directory.
rm -rf "${tmp}"
tmp=""

echo "Wrapper successfully generated and installed"
exit 0

